/*-------------------------------------------------------------------
2维泊松方程并行求解，默认格点数(65x65)
方程:
						- Laplacian u = 1      -1 < x, y < 1

边界：          
						上边界为1, 下边界为0, 左右两侧从0到1(从下边界至上边界是0到1)


Todo: 用DMDA来管理，可以访问全局X
https://petsc.org/release/src/ksp/ksp/tutorials/ex14f.F90.html
-------------------------------------------------------------------*/

#include <petscksp.h>
#include <petscmat.h>
#include <time.h>
#include <stdlib.h>

static char help[] = "Solves 2D Laplacian using matrix free.\n\n";
extern PetscErrorCode setVecBlockValue(Vec b, int gridNum);
extern PetscErrorCode poissonMult(Mat A, Vec x, Vec y);

typedef struct
{
    PetscReal   param;
    PetscInt    gridNum;

    PetscReal   x_start;
    PetscReal   x_end;
    PetscReal   y_start;
    PetscReal   y_end;

    PetscMPIInt rank;

    Vec         x;

} userData;


int main(int argc, char **args)
{
	PetscErrorCode 	ierr;
    KSP            	ksp;
    userData        user;
    PetscReal      	norm;
	Vec            	x, b, r;
	Mat				A;
	PetscInt		gridNum = 65, its;
    PetscInt        localSize = 0;
    PetscMPIInt     rank;
    PetscInt        plotflag = 0;
    time_t          time1, time2;

    // 初始化
    ierr = PetscInitialize(&argc, &args, (char*)0, help); if (ierr) return ierr;
    
    // 命令行参数
    ierr = PetscOptionsGetInt(NULL, NULL, "-gridnum", &gridNum, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetInt(NULL, NULL, "-plot", &plotflag, NULL); CHKERRQ(ierr);

    // 计时
    ierr = MPI_Comm_rank(PETSC_COMM_WORLD, &rank); CHKERRQ(ierr);
    if (!rank)
    {
        time1 = time(NULL);
    }

    // 格点数量
    PetscInt gridPointNum = gridNum*gridNum;
    user.param = 1.0;
    user.gridNum = gridNum;
    user.x_start = -1.0;
    user.x_end = 1.0;
    user.y_start = -1.0;
    user.y_end = 1.0;
    user.rank = rank;

    // b
    ierr = VecCreate(PETSC_COMM_WORLD, &b); CHKERRQ(ierr);
    ierr = VecSetSizes(b, PETSC_DECIDE, gridPointNum); CHKERRQ(ierr);
    ierr = VecSetFromOptions(b); CHKERRQ(ierr);
    ierr = VecDuplicate(b, &x); CHKERRQ(ierr);
    ierr = VecDuplicate(b, &r); CHKERRQ(ierr);
    ierr = VecDuplicate(b, &(user.x)); CHKERRQ(ierr);

    ierr = setVecBlockValue(b, gridNum); CHKERRQ(ierr);

    // A
    ierr = VecGetLocalSize(b, &localSize); CHKERRQ(ierr);
    ierr = MatCreateShell(PETSC_COMM_WORLD, localSize, localSize, gridPointNum, gridPointNum, (void *)&user, &A); CHKERRQ(ierr);
    ierr = MatShellSetOperation(A, MATOP_MULT, (void(*)(void))poissonMult); CHKERRQ(ierr);
    ierr = MatSetFromOptions(A); CHKERRQ(ierr);

    ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

    // ksp
    ierr = KSPCreate(PETSC_COMM_WORLD, &ksp); CHKERRQ(ierr);
    ierr = KSPSetType(ksp, KSPGMRES);
    ierr = KSPSetOperators(ksp, A, A); CHKERRQ(ierr);
    // ierr = KSPSetUp(ksp); CHKERRQ(ierr);

    // 设置ksp的迭代参数
    // ierr = KSPSetTolerances(ksp, 1.e-2/gridPointNum, 1.e-50, PETSC_DEFAULT, PETSC_DEFAULT); CHKERRQ(ierr);

    ierr = KSPSetFromOptions(ksp); CHKERRQ(ierr);

    // 求解
    ierr = KSPSolve(ksp, b, x); CHKERRQ(ierr);

    ierr = KSPGetIterationNumber(ksp, &its); CHKERRQ(ierr);
    PetscPrintf(PETSC_COMM_WORLD, "Iterations: %D\n", its);CHKERRQ(ierr);

    // 解算结果
    ierr = MatMult(A, x, r); CHKERRQ(ierr);
    ierr = VecAXPY(r, -1.0, b); CHKERRQ(ierr);
    ierr = VecNorm(r, NORM_2, &norm); CHKERRQ(ierr);

    PetscPrintf(PETSC_COMM_WORLD, "Residual norm %lf\n", (double)norm);
    PetscPrintf(PETSC_COMM_WORLD, "Average norm %lf\n", (double)norm / gridPointNum);

    // 输出结果到文件
    PetscViewer myViewer;
    ierr = PetscViewerASCIIOpen(PETSC_COMM_WORLD, "mpi_poisson2d.txt", &myViewer); CHKERRQ(ierr);
    ierr = VecView(x, myViewer); CHKERRQ(ierr);

    // 调用python脚本绘图
    if (!rank && plotflag)
    {
        system("./plot3d.py mpi_poisson2d.txt");
    }

    // 计算时间
    if (!rank)
    {
        time2 = time(NULL);
        printf("Time: %ld\n", time2 - time1);
    }

    // 释放内存
    // ierr = KSPDestroy(&ksp); CHKERRQ(ierr);
    ierr = VecDestroy(&x); CHKERRQ(ierr);
    ierr = VecDestroy(&b); CHKERRQ(ierr);
    ierr = MatDestroy(&A); CHKERRQ(ierr);
    ierr = VecDestroy(&r); CHKERRQ(ierr);

    ierr = PetscFinalize();
    return ierr;
}

PetscErrorCode setVecBlockValue(Vec b, int gridNum)
{
    PetscErrorCode      ierr;
    PetscInt            index, i, j, Istart, Iend, nloc;
    PetscScalar         *barray;

    PetscFunctionBeginUser;
    
    // 获取索引开始和结束
    ierr = VecGetOwnershipRange(b, &Istart, &Iend); CHKERRQ(ierr);
    ierr = VecGetArray(b, &barray); CHKERRQ(ierr);

    for (index = Istart; index < Iend; index++)
    {
        i = index % gridNum;
        j = index / gridNum;

        if (0 == j)
        {
            // y = -1
            barray[index - Istart] = 0.0;
        }
        else if (gridNum-1 == j)
        {
            // y = 1
            barray[index - Istart] = 1.0;
        }
        else if (0 == i || gridNum-1 == i)
        {
            barray[index - Istart] = j / (gridNum - 1.0);
        }
        else
        {
            barray[index - Istart] = 1.0;
        }
    }

    ierr = VecRestoreArray(b, &barray); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}


// 自定义矩阵乘法
// y = Ax
PetscErrorCode poissonMult(Mat A, Vec x, Vec y)
{
    PetscErrorCode      ierr;
    userData            *puser;
    const PetscScalar   *px;
    PetscScalar         *py;
    PetscInt            index, i, j, n, start, end;
    PetscScalar         dHxdHx, dHydHy, v[5];

    PetscFunctionBeginUser;
    MatShellGetContext(A, (void *)&puser);
    ierr = MatGetOwnershipRange(A, &start, &end); CHKERRQ(ierr);
    ierr = VecCopy(x, puser->x); CHKERRQ(ierr);
    ierr = VecGetArrayRead(puser->x, &px); CHKERRQ(ierr);
    ierr = VecGetArray(y, &py); CHKERRQ(ierr);

    // ierr = VecGetSize(x, &n); CHKERRQ(ierr);
    // ierr = VecGetLocalSize(puser->x, &n); CHKERRQ(ierr);
    // printf("rank start end n: %d %d  %d %d\n", puser->rank, start, end, n);
    // printf("rank: %d\n", puser->rank);

	dHxdHx = 1.0 / ((puser->x_end - puser->x_start) * (puser->x_end - puser->x_start) / ((PetscReal)((puser->gridNum-1) * (puser->gridNum-1))));
	dHydHy = 1.0 / ((puser->y_end - puser->y_start) * (puser->y_end - puser->y_start) / ((PetscReal)((puser->gridNum-1) * (puser->gridNum-1))));
    
    for (index = start; index < end; index++)
    {
        i = index % puser->gridNum;
        j = index / puser->gridNum;

        if (i == 0 || i == puser->gridNum-1 || j == 0 || j == puser->gridNum-1)
        {
            // 边界
            py[index - start] = px[index - start];
        }
        else
        {
            v[0] = -dHydHy * px[(j - 1) * puser->gridNum + i - start];
            v[1] = -dHxdHx * px[j * puser->gridNum + i - 1 - start];
            v[2] = 2.0*(dHxdHx + dHydHy) * px[j * puser->gridNum + i - start];
            v[3] = -dHxdHx * px[j * puser->gridNum + i + 1 - start];
            v[4] = -dHydHy * px[(j + 1) * puser->gridNum + i - start];
            
            py[index - start] = v[0] +  v[1] +  v[2] +  v[3] +  v[4];
        }
    }

    ierr = VecRestoreArrayRead(puser->x, &px); CHKERRQ(ierr);
    ierr = VecRestoreArray(y, &py); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}